package crypto

import (
	"bytes"

	"crypto/rand"
	"crypto/sha256"

	"golang.org/x/crypto/argon2"
	"golang.org/x/crypto/bcrypt"
	"golang.org/x/crypto/hkdf"
	"golang.org/x/crypto/pbkdf2"
	"golang.org/x/crypto/scrypt"
)

func NewPassword() *password {
	return &password{
		keyLen:         32,
		salt:           []byte("salt"),
		context:        []byte("context"),
		iterationCount: 1000,
		cost:           bcrypt.DefaultCost,
		time:           1,
		memory:         64 * 1024,
		threads:        4,
		n:              32768,
		r:              8,
		p:              1,
	}
}

type password struct {
	keyLen         int    // 密钥长度 Argon2 SCrypt Hkdf Pbkdf2
	salt           []byte // 盐值 Argon2 SCrypt Hkdf Pbkdf2
	context        []byte // 上下文 Hkdf
	iterationCount int    // 迭代次数 Pbkdf2
	cost           int    // 花费 BCrypt
	time           uint32 // 时间 Argon2
	memory         uint32 // 内存数 Argon2
	threads        uint8  // 线程数 Argon2
	n              int    // SCrypt
	r              int    // SCrypt
	p              int    // SCrypt
}

func (r *password) KeyLen(keyLen int) *password {
	r.keyLen = keyLen

	return r
}
func (r *password) Salt() *password {
	r.salt = make([]byte, 32)
	_, err := rand.Read(r.salt)
	if err != nil {
		panic(err)
	}

	return r
}
func (r *password) Context(context []byte) *password {
	r.context = context

	return r
}
func (r *password) IterationCount(iterationCount int) *password {
	r.iterationCount = iterationCount

	return r
}
func (r *password) Cost(cost int) *password { // 4-31|10
	if cost >= bcrypt.MinCost && cost <= bcrypt.MaxCost {
		r.cost = cost
	}

	return r
}

// Argon2i Argon2id
// keyLen
// salt
// time
// memory
// threads
func (r *password) Argon2i(source []byte) []byte {
	return argon2.Key(source, r.salt, r.time, r.memory, r.threads, uint32(r.keyLen))
}
func (r *password) Argon2id(source []byte) []byte {
	return argon2.IDKey(source, r.salt, r.time, r.memory, r.threads, uint32(r.keyLen))
}

// Hkdf
// keyLen
// salt
// context
func (r *password) Hkdf(source []byte) []byte {
	hash := make([]byte, r.keyLen)

	hasher := hkdf.New(sha256.New, source, r.salt, r.context)
	_, err := hasher.Read(hash)
	if err != nil {
		panic(err)
	}

	return hash
}

// Pbkdf2
// keyLen
// salt
// iterationCount
func (r *password) Pbkdf2(source []byte) []byte {
	return pbkdf2.Key(source, r.salt, r.iterationCount, r.keyLen, sha256.New)
}

// BCrypt
// cost 4-31|10
func (r *password) BCrypt(source []byte) []byte {
	hash, err := bcrypt.GenerateFromPassword(source, r.cost)
	if err != nil {
		panic(err)
	}

	return hash
}
func (r *password) BCryptVerify(source, hashed []byte) bool {
	err := bcrypt.CompareHashAndPassword(hashed, source)
	return err == nil
}

// SCrypt
// keyLen
// salt
// n
// r
// p
func (r *password) SCrypt(source []byte) ([]byte, []byte) {
	hash, err := scrypt.Key(source, r.salt, r.n, r.r, r.p, r.keyLen)
	if err != nil {
		panic(err)
	}

	return r.salt, hash
}
func (r *password) SCryptVerify(source, hashed []byte) bool {
	calculatedHash, err := scrypt.Key(source, r.salt, r.n, r.r, r.p, r.keyLen)
	if err != nil {
		return false
	}

	return bytes.Equal(hashed, calculatedHash)
}
